package com.shanzhi.dify.utils;

import com.shanzhi.dify.entity.AgentEntity;
import com.shanzhi.dify.entity.RetrievalModelRequestEntity;
import com.shanzhi.dify.entity.WorkFlowsRunEntity;

/**
 * 用于Dify的静态注入数据的工具类
 */
public class RebuildParamsUtils {

    /**
     * Agent模式不能使用Blocking 模式，只能使用非阻塞模式
     * 设置Agent模式为流式返回
     *
     * @param agentEntity
     * @return
     */
    public static void rebuildAgentEntity(AgentEntity agentEntity) {
        WorkFlowsRunEntity workFlowsRunEntity = agentEntity.getWorkFlowsRunEntity();
        workFlowsRunEntity.setResponse_mode("streaming");
    }



    /**
     * 创建带有固定参数的检索请求实体，只允许自定义查询内容
     *
     * @param query 查询字符串
     * @return 配置好的 RetrievalModelRequestEntity 对象
     */
    public static RetrievalModelRequestEntity createFixedRetrievalRequest(String query) {
        return RetrievalModelRequestEntity.builder()
                .query(query)
                .retrieval_model(RetrievalModelRequestEntity.RetrievalModel.builder()
                        .search_method("hybrid_search")
                        .reranking_enable(true)
                        .reranking_mode("reranking_model")
                        .reranking_model(RetrievalModelRequestEntity.RerankingModel.builder()
                                .reranking_provider_name("langgenius/tongyi/tongyi")
                                .reranking_model_name("gte-rerank")
                                .build())
                        .weights(RetrievalModelRequestEntity.Weights.builder()
                                .weight_type("customized")
                                .keyword_setting(RetrievalModelRequestEntity.KeywordSetting.builder()
                                        .keyword_weight(0.3)
                                        .build())
                                .vector_setting(RetrievalModelRequestEntity.VectorSetting.builder()
                                        .vector_weight(0.5)
                                        .embedding_model_name("text-embedding-v3")
                                        .embedding_provider_name("langgenius/tongyi/tongyi")
                                        .build())
                                .build())
                        .top_k(7)
                        .score_threshold_enabled(true)
                        .score_threshold(0.4)
                        .build())
                .build();
    }
}
